/*
* Copyright (C) 2020 ~ 2024 chaigec All Rights Reserved.
*
* Author:     chaigec <chaigec@163.com>
* Maintainer: chaigec <chaigec@163.com>

 * Use of this source code is governed by MIT license that can be found in the
 * LICENSE file in the root of the source tree. All contributing project authors
 * may be found in the AUTHORS file in the root of the source tree.
*/

#include "MkPsMuxer.h"
#include "MkPsUtil.h"
#include "MkUtil/MkLog.h"

#define MaxPesLen 0xFF80
Uint32 CMkPsMuxer::PacketEsPacket(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, const MkEsPacket& EsPacket, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, CMkBufferList& ReBufferList)
{
    Uint32 nCurUsed = 0;
    Uint8* pOffset = pBuf;
    //video
    if (MkCodecCavs > EsPacket.CodecType) {
        PacketPsHeader(VideoParam, AudioParam, EsPacket.Pts, pOffset, nLen, nCurUsed, ReBufferList);
        pOffset += nCurUsed;
        nLen -= nCurUsed;
        if (EsPacket.bKeyPacket) {
            PacketPsSystem(VideoParam, AudioParam, pOffset, nLen, nCurUsed, ReBufferList);
            pOffset += nCurUsed;
            nLen -= nCurUsed;
            PacketPsMap(VideoParam, AudioParam, pOffset, nLen, nCurUsed, ReBufferList);
            pOffset += nCurUsed;
            nLen -= nCurUsed;
        }
    }
    MkListBuffer::const_iterator it = EsPacket.BufferList.GetBufferList().begin();
    for (; it != EsPacket.BufferList.GetBufferList().end(); it++) {
        Uint32 nOffset = MaxPesLen;
        BOOL bAppendPs = FALSE;
        while (nOffset < it->GetLength()) {
            PacketPes(VideoParam, AudioParam, (const Uint8*)it->GetBuffer() + nOffset - MaxPesLen, MaxPesLen, pOffset, nLen, nCurUsed, EsPacket.CodecType, bAppendPs, EsPacket.Pts, EsPacket.Dts, ReBufferList);
            bAppendPs = TRUE;
            nOffset += MaxPesLen;
            pOffset += nCurUsed;
            nLen -= nCurUsed;
        }
        PacketPes(VideoParam, AudioParam, (const Uint8*)it->GetBuffer() + nOffset - MaxPesLen, it->GetLength() + MaxPesLen - nOffset, pOffset, nLen, nCurUsed, EsPacket.CodecType, bAppendPs, EsPacket.Pts, EsPacket.Dts, ReBufferList);
        pOffset += nCurUsed;
        nLen -= nCurUsed;
    }
    nUsed = pOffset - pBuf;
    return NoneError;
}

Uint32 CMkPsMuxer::PacketPsHeader(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, Uint32 Pts, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, CMkBufferList& ReBufferList)
{
    Uint16 SystemClockReferenceExtension = 0;
    Uint32 ProgramMuxRate = VideoParam.GetBitRate();
    Uint8 Reserved = 0x1F;
    Uint8 Marker = 1;
    Uint8 PackStuffingLenght = 0;
    Uint32 i = 0;
    pBuf[i++] = 0x00;
    pBuf[i++] = 0x00;
    pBuf[i++] = 0x01;
    pBuf[i++] = 0xBA;
    Uint64 SCR = (Uint64)Pts * (MkDefaultClockRate / 1000) + 32;
    pBuf[i] = 0x01 << 6;                           //2bit  always 0x01
    pBuf[i] |= (((SCR >> 30) & 0x07)) << 3;        //3bit  SCR 30-31
    pBuf[i] |= (Marker & 0x01) << 2;               //1 bit maker
    pBuf[i] |= ((SCR >> 28) & 0x03);               //2bit  SCR 28-29
    i++;

    pBuf[i++] = ((SCR >> 20) & 0xFF);             //8bit SCR 20-27

    pBuf[i] = ((SCR >> 15) & 0x1F) << 3;           //5bit SRC 15-19
    pBuf[i] |= (Marker & 0x01) << 2;               //1 bit maker
    pBuf[i] |= (SCR >> 13) & 0x03;                 //2bit SRC 13-14
    i++;

    pBuf[i++] = (SCR >> 5) & 0xFF;                 //8bit SRC 5-12

    pBuf[i] = (SCR & 0x1F) << 3;                   //5bit SRC 0-4
    pBuf[i] |= (Marker & 0x01) << 2;               //1bit marker
    pBuf[i] |= (SystemClockReferenceExtension >> 7) & 0x03; //2bit system clock reference extension 7-8
    i++;

    pBuf[i] = (SystemClockReferenceExtension & 0x7F) << 1; //7bit system clock reference extension 0-6
    pBuf[i] |= Marker & 0x01;                       //1bit marker
    i++;

    pBuf[i++] = (ProgramMuxRate >> 14) & 0xFF;     //8bit program mux rate 14-21
    pBuf[i++] = (ProgramMuxRate >> 6) & 0xFF;      //8bit program mux rate 6-13

    pBuf[i] = (ProgramMuxRate & 0x3F) << 2;        //6bit program mux rate 0-5
    pBuf[i] |= (Marker & 0x01) << 1;               //1bit marker
    pBuf[i] |= Marker & 0x01;;                     //1bit marker
    i++;

    pBuf[i] = (Reserved & 0x1F) << 3;         //5bit reserved
    pBuf[i] |= PackStuffingLenght & 0x07;   //3bit;
    i++;

    i += PackStuffingLenght;
    ReBufferList.Append(pBuf, i);
    nUsed = i;
    return NoneError;
}

Uint32 CMkPsMuxer::PacketPes(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, const Uint8* pData, Uint32 nDataLen, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, const MkCodecType& CodecType, BOOL bAppendPes, Uint32 Pts, Uint32 Dts, CMkBufferList& ReBufferList)
{
    BOOL bVideo = MkCodecCavs >= CodecType;
    Uint8  PesHeaderDataLength = 10;  //dts and pts 
    nUsed = PesHeaderDataLength + 6 + 3;  // 6 bytes is start 4 bytes and pes len 2 bytes  3 bytes is the pesheader param
    Uint16  PesPacketLen = nDataLen + PesHeaderDataLength + 3;
    if (bVideo && !bAppendPes) {
        nUsed += sizeof(Uint32);  //add nalu start code
        PesPacketLen += sizeof(Uint32);
    }
    if (!bVideo && MkCodecAac == CodecType) {
        nUsed += MkAacAdtsLen;  //add aac adts header 7 bytes
        PesPacketLen += MkAacAdtsLen;
    }
    if (nLen < nUsed) {
        MkErrorLog("input pbuf len:%d this has len:%d\n", nLen, nUsed);
        return NotHasEnoughBuffer;
    }
    Uint8   OriginalOrCopy = 0;
    Uint8   CopyRight = 0;
    Uint8   DataAlignmentIndicator = 1;
    Uint8   PesPriority = 1;
    Uint8   PesScramblingControl = 0;
    Uint8   FixBit = 2;

    Uint8   PesExtensionFlag = 0;
    Uint8   PesCrcFlag = 0;
    Uint8   AdditionalCopyInfoFlag = 0;
    Uint8   DsmTrickModeFlag = 0;
    Uint8   EsRateFlag = 0;
    Uint8   EscrFlag = 0;
    Uint8   PtsDtsFlags = 0x03;
    Uint8   Marker = 1;


    Uint32 i = 0;
    pBuf[i++] = 0x00;
    pBuf[i++] = 0x00;
    pBuf[i++] = 0x01;
    pBuf[i++] = bVideo ? 0xE0 : 0xC0;
    pBuf[i++] = (PesPacketLen >> 8) & 0xFF;
    pBuf[i++] = PesPacketLen & 0xFF;

    pBuf[i] = (FixBit & 0x03) << 6;                   //2bit
    pBuf[i] |= (PesScramblingControl & 0x01) << 4;   //2bit
    pBuf[i] |= (PesPriority & 0x01) << 3;            //1bit
    pBuf[i] |= (DataAlignmentIndicator & 0x01) << 2; //1bit
    pBuf[i] |= (CopyRight & 0x01) << 1;              //1bit
    pBuf[i] |= OriginalOrCopy & 0x01;                //1bit
    i++;

    pBuf[i] = (PtsDtsFlags & 0x03) << 6;            //2bit
    pBuf[i] |= (EscrFlag & 0x01) << 5;               //1bit
    pBuf[i] |= (EsRateFlag & 0x01) << 4;             //1bit
    pBuf[i] |= (DsmTrickModeFlag & 0x01) << 3;       //1bit
    pBuf[i] |= (AdditionalCopyInfoFlag & 0x01) << 2; //1bit
    pBuf[i] |= (PesCrcFlag & 0x01) << 1;             //1bit
    pBuf[i] |= PesExtensionFlag & 0x01;              //1bit
    i++;

    pBuf[i++] = PesHeaderDataLength & 0xFF;
    if (0x02 == (PtsDtsFlags & 0x02)) {
        //add pts and dts
        Uint64 PsPts = (Uint64)Pts * (MkDefaultClockRate / 1000) + 32;
        pBuf[i] = (PtsDtsFlags & 0x03) << 4;                           //4bit
        pBuf[i] |= (((PsPts >> 30) & 0x07)) << 1;      //3bit  PsPts 30-31
        pBuf[i] |= Marker & 0x01;                      //1bit
        i++;
        pBuf[i++] = (PsPts >> 22) & 0xFF;              //8bit PsPts 22-29
        pBuf[i] = ((PsPts >> 15) & 0x7F) << 1;         //7bit PsPts 15-21
        pBuf[i] |= Marker & 0x01;                      //1bit
        i++;

        pBuf[i++] = (PsPts >> 7) & 0xFF;               //8bit PsPts 7-14
        pBuf[i] = (PsPts & 0x7F) << 1;                //7bit PsPts 0-6
        pBuf[i] |= Marker & 0x01;                      //1bit
        i++;
    }
    if (0x01 == (PtsDtsFlags & 0x01)) {
        Uint64 PsDts = (Uint64)Dts * (MkDefaultClockRate / 1000);
        pBuf[i] = 0x01 << 4;                           //4bit
        pBuf[i] |= (((PsDts >> 30) & 0x07)) << 1;      //3bit  PsDts 30-31
        pBuf[i] |= Marker & 0x01;                      //1bit
        i++;
        pBuf[i++] = (PsDts >> 22) & 0xFF;              //8bit PsDts 22-29
        pBuf[i] = ((PsDts >> 15) & 0x7F) << 1;         //7bit PsDts 15-21
        pBuf[i] |= Marker & 0x01;                      //1bit
        i++;

        pBuf[i++] = (PsDts >> 7) & 0xFF;               //8bit PsDts 7-14
        pBuf[i] = (PsDts & 0x7F) << 1;                 //7bit PsDts 0-6
        pBuf[i] |= Marker & 0x01;                      //1bit
        i++;
    }

    if (bVideo && !bAppendPes) {
        //add video nalu start
        memcpy(pBuf + i, &NaluStartCode, sizeof(Uint32));
        i += sizeof(Uint32);
    }
    if (!bVideo && MkCodecAac == CodecType) {
        CMkMediaUtil::CreateAacAdts(AudioParam.GetObjectType(), AudioParam.GetSampleRate(), AudioParam.GetChannel(), nDataLen, &pBuf[i]);
        i += MkAacAdtsLen;
    }
    ReBufferList.Append(pBuf, i);
    ReBufferList.Append(pData, nDataLen);
    nUsed = i;
    return NoneError;
}

Uint32 CMkPsMuxer::PacketPsMap(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, CMkBufferList& ReBufferList)
{
    Uint16 ProgramStreamMapLength = 10;
    Uint8 CurrentNextIndicator = 0x01;
    Uint8 Reserved = 0xFF;
    Uint8 Marker = 1;
    Uint8 ProgramStreamMapVersion = 0x00;
    Uint16 ProgramStreamInfoLength = 0;
    Uint16 ElementaryStreamMapLength = 0x00;
    Uint32 Crc = 0xF4DCBD45;
    Uint32 i = 0;

    //has video
    if (MkCodecUnknown != VideoParam.GetCodec()) {
        ElementaryStreamMapLength += 4;
    }
    if (MkCodecUnknown != AudioParam.GetCodec()) {
        ElementaryStreamMapLength += 4;
    }
    ProgramStreamMapLength += ElementaryStreamMapLength;

    pBuf[i++] = 0x00;
    pBuf[i++] = 0x00;
    pBuf[i++] = 0x01;
    pBuf[i++] = 0xBC;

    pBuf[i++] = (ProgramStreamMapLength >> 8) & 0xFF;
    pBuf[i++] = ProgramStreamMapLength & 0xFF;
    pBuf[i] = (CurrentNextIndicator & 0x01) << 7;   //1bit
    pBuf[i] |= (Reserved & 0x03) << 5;              //2bit
    pBuf[i] |= ProgramStreamMapVersion & 0x1F;      //5bit
    i++;
    pBuf[i] = (Reserved & 0x7F) << 1;           //7bit
    pBuf[i] |= Marker & 0x01;                   //1bit
    i++;

    pBuf[i++] = (ProgramStreamInfoLength >> 8) & 0xFF;
    pBuf[i++] = ProgramStreamInfoLength & 0xFF;

    pBuf[i++] = (ElementaryStreamMapLength >> 8) & 0xFF;
    pBuf[i++] = ElementaryStreamMapLength & 0xFF;
    if (MkCodecUnknown != VideoParam.GetCodec()) {
        pBuf[i++] = CMkPsUtil::MkCodecTypeToPsCodec(VideoParam.GetCodec());
        pBuf[i++] = 0xE0;
        pBuf[i++] = 0;
        pBuf[i++] = 0;
    }
    if (MkCodecUnknown != AudioParam.GetCodec()) {
        pBuf[i++] = CMkPsUtil::MkCodecTypeToPsCodec(AudioParam.GetCodec());
        pBuf[i++] = 0xC0;
        pBuf[i++] = 0;
        pBuf[i++] = 0;
    }
    pBuf[i++] = (Crc >> 24) & 0xFF;
    pBuf[i++] = (Crc >> 16) & 0xFF;
    pBuf[i++] = (Crc >> 8) & 0xFF;
    pBuf[i++] = Crc & 0xFF;
    ReBufferList.Append(pBuf, i);
    nUsed = i;
    return NoneError;
}

Uint32 CMkPsMuxer::PacketPsSystem(const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, Uint8* pBuf, Uint32 nLen, Uint32& nUsed, CMkBufferList& ReBufferList)
{
    Uint8 Maker = 1;
    Uint16 HeaderLen = PsSystemHeaderLen - 6;
    Uint32 RateBound = VideoParam.GetBitRate();
    Uint8 AudioBound = 1;
    Uint8 FixedFlag = 0;
    Uint8 CspsFlag = 0;
    Uint8 SysVideoLockFlag = MkCodecUnknown == VideoParam.GetCodec() ? 0 : 1;
    Uint8 SysAudioLockFlag = MkCodecUnknown == AudioParam.GetCodec() ? 0 : 1;
    Uint8 VideoBound = 1;
    Uint8 PacketRateRestrictionFlag = 0;
    Uint8 Reserved = 0x7F;

    Uint32 i = 0;
    pBuf[i++] = 0x00;
    pBuf[i++] = 0x00;
    pBuf[i++] = 0x01;
    pBuf[i++] = 0xBB;
    pBuf[i++] = (HeaderLen >> 8) & 0xFF;
    pBuf[i++] = HeaderLen & 0xFF;
    pBuf[i] = Maker << 7;
    pBuf[i] |= (RateBound >> 15) & 0x7F;  //7bit rate bound 15-21
    i++;

    pBuf[i++] = (RateBound >> 7) & 0xFF;  //8bit rate bound 7-14

    pBuf[i] = (RateBound & 0x1F) << 1;   //7bit rate bond 0-6
    pBuf[i] |= Maker;
    i++;

    pBuf[i] = (AudioBound & 0x3F) << 2;  //6bit
    pBuf[i] |= (FixedFlag & 0x01) << 1;  //1bit
    pBuf[i] |= CspsFlag & 0x01;          //1bit
    i++;

    pBuf[i] = (SysAudioLockFlag & 0x01) << 7;       //1bit
    pBuf[i] |= (SysVideoLockFlag & 0x01) << 6;      //1bit
    pBuf[i] |= (Maker & 0x01) << 5;                 //1bit
    pBuf[i] |= (VideoBound & 0x1F);                 //5bit
    i++;

    pBuf[i] = (PacketRateRestrictionFlag & 0x01) << 7;  //1bit
    pBuf[i] |= (Reserved & 0x7F);                       //7bit
    i++;

    ReBufferList.Append(pBuf, i);
    nUsed = i;
    return NoneError;
}
